#!/usr/bin/env ruby

# Copyright (c) 2010-2012 Microsoft Corp.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#	http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

require 'rubygems'
require 'ruby-aws'
require 'digest/md5'

require './include/config'
require './include/proxy'
require './include/requester'
require './include/templates'

@config = MOS::Config.create
@mturk  = Amazon::WebServices::MechanicalTurkRequester.new :Host => @config[:Server]

def getHITUrl(hitTypeId)
	if @mturk.host =~ /sandbox/
		"http://workersandbox.mturk.com/mturk/preview?groupId=#{hitTypeId}" # Sandbox Url
	else
		"http://mturk.com/mturk/preview?groupId=#{hitTypeId}" # Production Url
	end
end

def factors(n)
	return [] if n == 1
	factor = (2..n).to_a.delete_if {|x| n % x != 0}
end     

def createHITs
	if !File.exists?(@config[:HITInputFile])
		$stderr.printf("Error: HIT input file (%s) does not exist.\n", @config[:HITInputFile])
		$stderr.printf("Exiting...\n")
		exit
	end

	if !File.exists?(@config[:HITPropertiesFile])
		$stderr.printf("Error: HIT property file (%s) does not exist.\n", @config[:HITPropertiesFile])
		$stderr.printf("Exiting...\n")
		exit
	end

	input = Amazon::Util::DataReader.load(@config[:HITInputFile], :Tabular)
	props = Amazon::Util::DataReader.load(@config[:HITPropertiesFile], :Properties)

	input.each do |line|
		line[:sentence]   = line[:sentence].to_s
		line[:algorithm]  = line[:algorithm].to_s
		line[:identifier] = line[:sentence] + '_' + line[:algorithm]
	end
	input.sort!{ |a,b| a[:identifier] <=> b[:identifier] }
	input.each_index { |i| input[i][:identifier] = 'ID' + i.to_s } if @config[:RewriteIdentifier]

	input.each { |line| line[:URL] = CGI.escape(line[:URL]) } if @options[:external_question]

	algorithms = input.map { |line| line[:algorithm] }.uniq
	sentences  = input.map { |line| line[:sentence]  }.uniq

	input_by_sentence = MOS.indexByNonUniqueKey(input, :sentence)
	files_per_sentence = input_by_sentence.values.map { |file_list| file_list.length }

	if files_per_sentence.min != files_per_sentence.max
		$stderr.printf("Error: the set of sentences must be the same for every algorithm.\n")
		$stderr.printf("Exiting...\n")
		exit
	end

	dataset = []
	case @options[:design]
		when :mushra
			# multiple stimulus with hidden reference and anchor
			input_by_sentence.each_value do |i|
				ref_file = i.find { |file| file[:algorithm].match(/^Ref/i) }

				if ref_file == nil
					$stderr.printf("Error: at least one sentence does not have an associated reference.\n")
					$stderr.printf("Exiting...\n")
					exit
				end

				hit_data = MOS.buildInputDataSet(i)
				hit_data[:reference_URL] = ref_file[:URL]
				dataset << hit_data
			end
			@config[:SamplesPerHIT] = algorithms.length

		when :dsis
			# double stimulus with the reference first, no anchor
			input_by_sentence.each_value do |i|
				ref_file = i.find { |file| file[:algorithm].match(/^Ref/i) }

				if ref_file == nil
					$stderr.printf("Error: at least one sentence does not have an associated reference.\n")
					$stderr.printf("Exiting...\n")
					exit
				end

				while i.length > 1
					test_file = i.find { |file| !file[:algorithm].match(/^Ref/i) }

					if test_file == nil
						$stderr.printf("Error: could not find test file matching the reference.\n")
						$stderr.printf("Exiting...\n")
					end

					i.delete(test_file)

					hit_data = MOS.buildInputDataSet([test_file])
					hit_data[:reference_URL] = ref_file[:URL]
					dataset << hit_data
				end
			end
			@config[:SamplesPerHIT] = 1

		when :same_sentence
			input_by_sentence.each_value do |i|
				dataset << MOS.buildInputDataSet(i)
			end
			@config[:SamplesPerHIT] = algorithms.length

		when :random_sentences
			input.shuffle!

			printf("Please enter number of files per HIT: ")
			@config[:SamplesPerHIT] = STDIN.gets.chomp.to_i

			while input.length > 0
				dataset << MOS.buildInputDataSet(input.pop(@config[:SamplesPerHIT]))
			end

		when :diff_sentences
			input_by_sentence = input_by_sentence.values

			f = factors(sentences.length)
			printf("Please enter number of files per HIT (must be one of the following: " + f.join(' ') + "): ")
			@config[:SamplesPerHIT] = STDIN.gets.chomp.to_i
			if f.find { |x| x == @config[:SamplesPerHIT] } == nil
				$stderr.printf("Error: the specified number of files per HIT must belong to the list above.\n");
				$stderr.printf("Exiting...\n")
				exit
			end

			fail unless (input_by_sentence.length % @config[:SamplesPerHIT]) == 0
			while input_by_sentence.length > 0 do
				input_by_sentence.shuffle!
				input_by_sentence.sort! { |x,y| y.length <=> x.length }
				input_by_sentence.each  { |x| x.shuffle! }

				dataset_item = []
				[@config[:SamplesPerHIT], input_by_sentence.length].min.times do |i|
					dataset_item << input_by_sentence[i].pop
				end
				dataset << MOS.buildInputDataSet(dataset_item)
				input_by_sentence.delete_if { |x| x.length == 0 }
			end
	end
	dataset.shuffle!

	@config.save

	if @options[:external_question]
		quest    = MOS.buildQuestionTemplateExternal(@options[:external_question], @config[:SamplesPerHIT])
	else
		if !File.exists?(@config[:HITQuestionFile] + '_header')
			$stderr.printf("Error: HIT question header (%s) does not exist.\n", @config[:HITQuestionFile] + '_header')
			$stderr.printf("Exiting...\n")
			exit
		end

		if !File.exists?(@config[:HITQuestionFile] + '_template')
			$stderr.printf("Error: HIT question template (%s) does not exist.\n", @config[:HITQuestionFile] + '_template')
			$stderr.printf("Exiting...\n")
			exit
		end

		if !File.exists?(@config[:HITQuestionFile] + '_footer')
			$stderr.printf("Error: HIT question footer (%s) does not exist.\n", @config[:HITQuestionFile] + '_footer')
			$stderr.printf("Exiting...\n")
			exit
		end

		header   = File.read(@config[:HITQuestionFile] + '_header')
		template = File.read(@config[:HITQuestionFile] + '_template')
		footer   = File.read(@config[:HITQuestionFile] + '_footer')
		quest    = MOS.buildQuestionTemplate(header, template, footer, @config[:SamplesPerHIT])
	end

	total_cost = props[:Reward][:Amount] * dataset.size * props[:MaxAssignments]
	total_cost *= 1.10	# Amazon's commission is 10%
	available_funds = @mturk.availableFunds

	if available_funds < total_cost
		$stderr.printf("Not enough funds: operation requires %.2f USD, account contains %.2f USD.\n",
				total_cost, available_funds)
		exit
	end

	printf("Create %d HITs with %d assignments each for a base cost of %.2f %s (server: %s)? (y/n) ",
			dataset.size, props[:MaxAssignments], total_cost, props[:Reward][:CurrencyCode],
			@config[:Server])
	response = STDIN.gets.chomp.downcase
	if response != 'y'
		printf("Exiting...\n")
		exit
	end

	puts("Creating HITs...")
	hits = @mturk.createHITs(props, quest, dataset)
	puts("Done.")

	hit_ids = hits[:Created].collect { |h| h[:HITId] }
	hit_type_id = hits[:Created].first[:HITTypeId]

	puts("Created HITs: #{hit_ids.join(' ')}")
	puts("Url: #{getHITUrl(hit_type_id)}")

	Amazon::Util::DataReader.save(@config[:HITSuccessFile], hits[:Created], :Tabular)
	Amazon::Util::DataReader.save(@config[:HITFailureFile], hits[:Failed],  :Tabular)
end

@options = { }
optparse = OptionParser.new do |opts|
	opts.banner = "Usage: #{$0} [options]"

	@options[:external_question] = nil
	opts.on('--external-question URL',
		'Generate HITs for an externally hosted question, accessible at URL' ) do |url|
		@options[:external_question] = url
	end

	# mushra: each HIT compares the same sentence, with a hidden anchor and a hidden reference
	# dsis: double stimulus impairment scale (DMOS where the first is always the reference)
	# same_sentence: each HIT compares the same sentence processed by all algorithms
	# diff_sentences: each HIT compares a random selection of files, never featuring a sentence twice
	# random_sentences: each HIT compares a random selection of files, without additional constraints
	@options[:design] = nil
	opts.on('--design DESIGN', [:mushra, :dsis, :same_sentence, :diff_sentences, :random_sentences],
			'Select test design (mushra, dsis, same_sentence, diff_sentences, random_sentences).') do |t|
		@options[:design] = t
	end

	opts.on('-h', '--help', 'Display this screen') do
		puts opts
		exit
	end

	opts.on_tail('')
	opts.on_tail('This script submits HITs by instantiating the HIT template with input parameters.')
end

begin
	optparse.parse!
rescue OptionParser::ParseError => e
	$stderr.puts("Error parsing command-line arguments: " + e.message)
	$stderr.puts("Try '#{$0} --help' for more information.")
	exit
end

if !@options[:design]
	$stderr.printf("No test design informed (use the --design switch).\n")
	$stderr.puts("Try '#{$0} --help' for more information.")
	exit
end

createHITs

